package com.codegen.controller;

import lombok.Data;
import org.springframework.web.bind.annotation.*;

import java.sql.*;
import java.util.ArrayList;
import java.util.List;

@RestController
@RequestMapping("/api/datasource")
public class DataSourceController {

    @PostMapping("/tables")
    public List<TableInfo> getTables(@RequestBody DataSourceInfo dataSourceInfo) throws Exception {
        List<TableInfo> tables = new ArrayList<>();
        String dbUrl = dataSourceInfo.getUrl();
        String dbName = dbUrl.substring(dbUrl.lastIndexOf("/") + 1, dbUrl.contains("?") ? dbUrl.indexOf("?") : dbUrl.length());
        try (Connection conn = DriverManager.getConnection(
                dataSourceInfo.getUrl(), dataSourceInfo.getUsername(), dataSourceInfo.getPassword())) {
            // 查询表基础信息
            String sql = "SELECT table_name, table_comment, engine, table_collation, auto_increment FROM information_schema.tables WHERE table_schema = ?";
            try (PreparedStatement ps = conn.prepareStatement(sql)) {
                ps.setString(1, dbName);
                try (ResultSet rs = ps.executeQuery()) {
                    while (rs.next()) {
                        TableInfo info = new TableInfo();
                        info.setName(rs.getString("table_name"));
                        info.setDescription(rs.getString("table_comment"));
                        info.setEngine(rs.getString("engine"));
                        info.setCollation(rs.getString("table_collation"));
                        // 获取字符集
                        info.setCharset(getCharsetByCollation(conn, rs.getString("table_collation")));
                        // 获取自增字段名
                        String autoIncCol = getAutoIncrementColumn(conn, dbName, rs.getString("table_name"));
                        info.setAutoIncrement(autoIncCol != null ? autoIncCol : "否");
                        tables.add(info);
                    }
                }
            }
        }
        return tables;
    }

    private String getCharsetByCollation(Connection conn, String collation) throws SQLException {
        if (collation == null) return "";
        String sql = "SELECT character_set_name FROM information_schema.collations WHERE collation_name = ?";
        try (PreparedStatement ps = conn.prepareStatement(sql)) {
            ps.setString(1, collation);
            try (ResultSet rs = ps.executeQuery()) {
                if (rs.next()) {
                    return rs.getString("character_set_name");
                }
            }
        }
        return "";
    }

    private String getAutoIncrementColumn(Connection conn, String dbName, String tableName) throws SQLException {
        String sql = "SELECT column_name FROM information_schema.columns WHERE table_schema = ? AND table_name = ? AND extra = 'auto_increment'";
        try (PreparedStatement ps = conn.prepareStatement(sql)) {
            ps.setString(1, dbName);
            ps.setString(2, tableName);
            try (ResultSet rs = ps.executeQuery()) {
                if (rs.next()) {
                    return rs.getString("column_name");
                }
            }
        }
        return null;
    }

    @PostMapping("/columns")
    public List<ColumnInfo> getColumns(@RequestBody ColumnRequest request) throws Exception {
        List<ColumnInfo> columns = new ArrayList<>();
        String dbUrl = request.getUrl();
        String dbName = dbUrl.substring(dbUrl.lastIndexOf("/") + 1, dbUrl.contains("?") ? dbUrl.indexOf("?") : dbUrl.length());
        try (Connection conn = DriverManager.getConnection(request.getUrl(), request.getUsername(), request.getPassword())) {
            String sql = "SELECT column_name, column_type, column_key, is_nullable, column_comment, extra FROM information_schema.columns WHERE table_schema = ? AND table_name = ? ORDER BY ordinal_position";
            try (PreparedStatement ps = conn.prepareStatement(sql)) {
                ps.setString(1, dbName);
                ps.setString(2, request.getTableName());
                try (ResultSet rs = ps.executeQuery()) {
                    while (rs.next()) {
                        ColumnInfo info = new ColumnInfo();
                        info.setName(rs.getString("column_name"));
                        info.setType(rs.getString("column_type"));
                        info.setPrimaryKey("PRI".equals(rs.getString("column_key")));
                        info.setNullable("YES".equals(rs.getString("is_nullable")));
                        info.setComment(rs.getString("column_comment"));
                        info.setAutoIncrement(rs.getString("extra") != null && rs.getString("extra").contains("auto_increment"));
                        columns.add(info);
                    }
                }
            }
        }
        return columns;
    }

    @Data
    public static class DataSourceInfo {
        private String url;
        private String username;
        private String password;
    }

    @Data
    public static class TableInfo {
        private String name;
        private String description;
        private String engine;
        private String autoIncrement;
        private String charset;
        private String collation;
    }

    @Data
    public static class ColumnRequest extends DataSourceInfo {
        private String tableName;
    }

    @Data
    public static class ColumnInfo {
        private String name;
        private String type;
        private boolean primaryKey;
        private boolean nullable;
        private String comment;
        private boolean autoIncrement;
    }
} 